import { PGVectorStore } from '@langchain/community/vectorstores/pgvector';
import { OllamaEmbeddings, ChatOllama } from '@langchain/ollama';
import { ChatPromptTemplate } from '@langchain/core/prompts';
import { RunnableSequence } from '@langchain/core/runnables';
import { StringOutputParser } from '@langchain/core/output_parsers';
import { Document } from '@langchain/core/documents';
import { MemoryVectorStore } from 'langchain/vectorstores/memory';
import { RecursiveCharacterTextSplitter } from 'langchain/text_splitter';
import { v4 as uuidv4 } from 'uuid';

// 定义会话数据结构
interface SessionData {
  tempRetriever: ReturnType<MemoryVectorStore['asRetriever']>; // 存储临时检索器
  timestamp: number;  // 最后活动时间戳，用于过期清理
}
// 全局会话存储 Map
const sessionStore = new Map<string, SessionData>();
const SESSION_TTL = 30 * 60 * 1000; // 会话过期时间：30分钟

// 定期清理过期会话 (每10分钟执行一次)
setInterval(() => {
  const now = Date.now();
  sessionStore.forEach((data, sessionId) => {
    if (now - data.timestamp > SESSION_TTL) {
      sessionStore.delete(sessionId);
      console.log(`⏰ 清理过期会话: ${sessionId}`);
    }
  });
}, 10 * 60 * 1000);

// 声明一个变量用于缓存检索器
let vectorStoreRetriever: Awaited<ReturnType<typeof initializeRetriever>> | null = null;

const ollama = new ChatOllama({
  baseUrl: 'http://192.168.31.159:11434',
  model: 'qwen:32b',
});

// 定义初始化检索器的函数
async function initializeRetriever() {
  console.log('🔧 正在初始化向量存储检索器...');
  const vectorStore = await PGVectorStore.initialize(
    new OllamaEmbeddings({
      baseUrl: 'http://192.168.31.159:11434',
      model: 'quentinz/bge-large-zh-v1.5:latest',
    }),
    {
      tableName: 'embedding_zqw_1024',
      columns: {
        idColumnName: 'embedding_id',
        vectorColumnName: 'embedding',
        contentColumnName: 'text',
        metadataColumnName: 'metadata',
      },
      postgresConnectionOptions: {
        host: '118.178.106.254',
        port: 15432,
        user: 'postgres',
        password: '6a29c7071ae1f1b9',
        database: 'postgres',
      },
    }
  );
  return vectorStore.asRetriever(3);
}

// 立即执行初始化函数并将 Promise 保存起来
const retrieverPromise = initializeRetriever();

async function getVectorStoreRetriever() {
  if (!vectorStoreRetriever) {
    vectorStoreRetriever = await retrieverPromise;
  }
  return vectorStoreRetriever;
}

// 创建 RAG Chain - 使用已检索的文档
async function createRAGChainWithDocs(retrievedDocs: Document[]) {
  // 定义主提示模板
  const TEMPLATE = `你是一个专业的知识问答助手，擅长基于提供的文档内容进行精准回答。

请严格按照以下格式回答问题：

**思考过程：**
[在这里分析检索到的文档片段，说明你是如何从这些片段中找到答案的，包括：
- 哪些片段与问题相关
- 关键信息在哪里
- 如何得出结论]

**最终答案：**
[给出简洁明确的答案]

相关文档内容：
{context}

注意：
1. 必须包含"思考过程："和"最终答案："两个部分
2. 思考过程要详细，要引用具体的文档片段
3. 如果文档中没有相关信息，在思考过程中说明，并在最终答案中告知用户`;

  const chatPrompt = ChatPromptTemplate.fromMessages([
    ['system', TEMPLATE],
    ['human', '{question}'],
  ]);

  // 🔧 关键修复：直接使用传入的已检索文档，而不是重新检索
  const context = retrievedDocs.map((doc) => doc.pageContent).join('\n\n');

  // 组装 RAG 链 - 简化版，直接使用已有的上下文
  const ragChain = RunnableSequence.from([
    // 将问题和上下文组合
    (input: { question: string }) => ({
      question: input.question,
      context: context, // 使用已检索的文档
    }),
    chatPrompt,
    ollama,
    new StringOutputParser(),
  ]);

  return ragChain;
}

// 处理上传文档的向量检索
async function createTempVectorStore(documentContent: string, filename: string) {
  // 按逻辑顺序进行切割
  const splitter = new RecursiveCharacterTextSplitter({
    chunkSize: 150, // 每个块的最大字符数
    chunkOverlap: 20, // 块之间的重叠字符数（增加重叠，避免关键信息丢失）
    separators: [
      '\n\n', // 优先按段落分割（最强语义边界）
      '\n', // 然后按换行分割
      '。', '！', '？', // 中文句子结尾（次强语义边界）
      '，', '；', '：', // 中文短语分隔符（弱语义边界）
      ' ', '', // 最后按空格或无分隔符切割（兜底）
    ]
  });

  const doc = new Document({
    pageContent: documentContent,
    metadata: { source: filename, type: 'uploaded' },
  });

  const splitDocs = await splitter.splitDocuments([doc]);
  console.log(`📄 文档切分完成，共 ${splitDocs.length} 个片段`);

  const embeddings = new OllamaEmbeddings({
    baseUrl: 'http://192.168.31.159:11434',
    model: 'quentinz/bge-large-zh-v1.5:latest',
  });

  const vectorStore = await MemoryVectorStore.fromDocuments(splitDocs, embeddings);
  return vectorStore.asRetriever(3);
}

// RAG 接口：使用 Vercel AI SDK 标准流式格式
export async function POST(req: Request) {
  try {
    // 1. 先处理header_Id
    let sessionId = req.headers.get('X-Session-ID');
    console.log(sessionId)
    
    if (!sessionId) sessionId = uuidv4();//如果没有则生成一个id
    const { messages, uploadedDocument } = await req.json();

    console.log('📨 RAG 模式收到请求', uploadedDocument ? '(包含上传文档)' : '');

    // 从最后一条消息中提取问题
    const lastMessage = messages[messages.length - 1];
    const question = lastMessage.parts?.find((p: { type: string; text?: string }) => p.type === 'text')?.text || lastMessage.content;

    let retrievedDocs: Document[] = [];

    // 2. 获取或创建临时检索器
    let sessionData = sessionStore.get(sessionId);

    // 如果有上传的新文档，或者会话中没有检索器，则创建并更新
    if (uploadedDocument && uploadedDocument.content) {
      console.log('🔍 正在创建新的临时文档检索器...');
      const tempRetriever = await createTempVectorStore(uploadedDocument.content, uploadedDocument.filename);
      // 更新或创建会话数据
      sessionStore.set(sessionId, {
        tempRetriever: tempRetriever,
        timestamp: Date.now()
      });
      sessionData = { tempRetriever, timestamp: Date.now() };
    }
    // 如果会话中存在检索器，则直接使用
    else if (sessionData && sessionData.tempRetriever) {
      console.log('🔍 使用会话中缓存的临时文档检索器...');
      sessionData.timestamp = Date.now(); // 更新活动时间
    }

    if (sessionData && sessionData.tempRetriever) {
      retrievedDocs = await sessionData.tempRetriever.invoke(question);

      // 如果临时文档中没有找到，再回退到数据库
      if (retrievedDocs.length === 0 || !retrievedDocs.some(doc => doc.pageContent.trim().length > 10)) {
        console.log('⚠️ 缓存的上传文档中未找到相关内容，转向数据库检索...');
        const dbRetriever = await getVectorStoreRetriever();
        retrievedDocs = await dbRetriever.invoke(question);
      }
    } else {
      // 如果没有任何临时检索器，则直接从数据库检索
      console.log('🔍 从数据库检索...');
      const retriever = await getVectorStoreRetriever();
      retrievedDocs = await retriever.invoke(question);
    }

    // ⚠️ 关键改进：检查检索结果的质量
    const hasValidResults = retrievedDocs.length > 0 && retrievedDocs.some(doc => doc.pageContent.trim().length > 20);
    
    // 如果没有找到有效的检索结果，转发到 normal-chat
    if (!hasValidResults) {
      console.log('⚠️ RAG 未找到相关内容，转发到 normal-chat...');
      
      // 转发到 normal-chat
      const normalChatResponse = await fetch(`${process.env.NEXT_PUBLIC_BASE_URL || 'http://localhost:3000'}/api/manage/normal-chat`, {
        method: 'POST',
        headers: { 'Content-Type': 'application/json' },
        body: JSON.stringify({ messages }),
      });

      if (!normalChatResponse.ok) {
        throw new Error('转发到 normal-chat 失败');
      }

      // 在返回前添加提示信息
      const encoder = new TextEncoder();
      const readableStream = new ReadableStream({
        async start(controller) {
          // 先输出提示信息
          const notice = '💡 提示：知识库中未找到相关信息，以下是基于通用知识的回答：\n\n';
          const data = `0:${JSON.stringify(notice)}\n`;
          controller.enqueue(encoder.encode(data));

          // 然后转发 normal-chat 的流式响应
          const reader = normalChatResponse.body?.getReader();
          if (reader) {
            while (true) {
              const { done, value } = await reader.read();
              if (done) break;
              controller.enqueue(value);
            }
          }
          controller.close();
        },
      });

      const responseHeaders = new Headers({
        'Content-Type': 'text/plain; charset=utf-8',
        'X-Vercel-AI-Data-Stream': 'v1',
        'X-Session-ID': sessionId,
        'X-Fallback-Mode': 'normal', // 标记已回退到 normal 模式
      });

      return new Response(readableStream, { headers: responseHeaders });
    }

    // 构建检索到的文本片段展示
    const retrievalInfo = retrievedDocs.map((doc: Document, index: number) => {
      const content = doc.pageContent.substring(0, 150);
      const source = doc.metadata?.source || '未知来源';
      const sourceType = doc.metadata?.type === 'uploaded' ? '📤 上传文档' : '💾 数据库';
      return `📄 片段${index + 1} (${sourceType}: ${source}):\n"${content}${doc.pageContent.length > 150 ? '...' : ''}"`;
    }).join('\n\n');
    
    // 🔧 关键修复：使用已检索的文档创建 RAG 链
    const ragChain = await createRAGChainWithDocs(retrievedDocs);
    const stream = await ragChain.stream({ question });

    // 使用 Vercel AI SDK 的标准流式格式
    const encoder = new TextEncoder();
    const readableStream = new ReadableStream({
      async start(controller) {
        try {
          // 模拟流式输出检索信息
          const retrievalPrefix = `🔍 检索到 ${retrievedDocs.length} 个相关文档片段：\n\n${retrievalInfo}\n\n---\n\n`;

          // 将检索信息分块流式输出（模拟打字效果）
          const chunkSize = 50; // 每次输出50个字符
          for (let i = 0; i < retrievalPrefix.length; i += chunkSize) {
            const chunk = retrievalPrefix.slice(i, i + chunkSize);
            const data = `0:${JSON.stringify(chunk)}\n`;
            controller.enqueue(encoder.encode(data));
            // 添加小延迟，让流式效果更明显（可选）
            await new Promise(resolve => setTimeout(resolve, 10));
          }

          // 然后输出模型的流式回答（包含思考过程和最终答案）
          for await (const chunk of stream) {
            const text = typeof chunk === 'string' ? chunk : String(chunk);
            // Vercel AI SDK 标准格式：0:"文本块"\n
            const data = `0:${JSON.stringify(text)}\n`;
            controller.enqueue(encoder.encode(data));
          }
          controller.close();
        } catch (error) {
          console.error('RAG 流式输出错误:', error);
          controller.error(error);
        }
      },
    });

    const responseHeaders = new Headers({
      'Content-Type': 'text/plain; charset=utf-8',
      'X-Vercel-AI-Data-Stream': 'v1',
      'X-Session-ID': sessionId, // 返回 Session ID 给前端
    });

    return new Response(readableStream, { headers: responseHeaders });
  } catch (error) {
    console.error('RAG 处理失败:', error);
    return new Response(JSON.stringify({ error: 'RAG 处理失败' }), {
      status: 500,
      headers: { 'Content-Type': 'application/json' },
    });
  }
}
